import os

from sentence_transformers import SentenceTransformer, CrossEncoder
import torch
import torchvision.models as models

from zhanshop.app import App
from zhanshop.env import Env


class ModelProvider():
    def getSentenceTransformerModelPath(self, modelName, saveDir='/runtime/model/'):
        """
        获取SentenceTransformer【文本语义】模型路径
        :param modelName:
        :param saveDir:
        :return:
        """
        saveDir = App.make(Env).rootPath + saveDir
        os.makedirs(saveDir, exist_ok=True)
        model_path = os.path.join(saveDir, modelName.replace("/", "__").lower())
        #print(model_path)
        # 如果模型不存在，则下载
        if not os.path.exists(model_path):
            model = SentenceTransformer(modelName)
            model.save(model_path)
        return str(model_path)

    def getCrossEncoderModelPath(self, modelName, saveDir='/runtime/model/'):
        """
        模型
        :param modelName:
        :param saveDir:
        :return:
        """
        saveDir = App.make(Env).rootPath + saveDir
        os.makedirs(saveDir, exist_ok=True)
        model_path = os.path.join(saveDir, modelName.replace("/", "__").lower())
        # print(model_path)
        # 如果模型不存在，则下载
        if not os.path.exists(model_path):
            model = CrossEncoder(modelName)
            model.save(model_path)
        return str(model_path)
    def getResNet50ModelPath(self, modelName, saveDir='/runtime/model/'):
        """
        获取ResNet50【深度卷积神经网络】模型路径
        :param modelName:
        :param saveDir:
        :return:
        """
        modelPath = App.make(Env).rootPath + saveDir + str(modelName).replace(".", "__").lower()+".pth"
        #model_path="./models/resnet50.pth"
        # 如果模型不存在，则下载并保存
        if not os.path.exists(modelPath):
            # print("下载ResNet50模型权重...")
            model = models.resnet50(weights=modelName)
            torch.save(model.state_dict(), modelPath)
            # print(f"模型权重已保存到: {model_path}")
        return str(modelPath)